import matplotlib.pyplot as plt
import numpy as np
font = {'family': 'Arial', 'size' : 16}
plt.rc('font', **font)

plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.it'] = 'Arial:italic'
plt.rcParams['mathtext.rm'] = 'Arial'
plt.rcParams['pdf.fonttype'] = 42

# Fixing random state for reproducibility
np.random.seed(19680801)

no_noise_data = [0.44659, 0.40085, 0.44659, 0.40085]
# linear noise
data = np.zeros((4, 3, 3))
# 1by1
data[0][0] = [0.46893, 0.46698, 0.47088] # low noise
data[0][1] = [0.47317, 0.49119, 0.50714] # medium noise
data[0][2] = [0.52819, 0.51975, 0.52764] # high noise
# onion
data[1][0] = [0.40654, 0.41613, 0.43156] # low noise
data[1][1] = [0.47442, 0.47987, 0.46687] # medium noise
data[1][2] = [0.58012, 0.56189, 0.56181] # high noise

# normal noise
# 1by1
data[2][0] = [0.5202, 0.52142, 0.51633] # low noise
data[2][1] = [0.59197, 0.57579, 0.58015] # medium noise
data[2][2] = [0.71255, 0.69392, 0.68804] # high noise
# onion
data[3][0] = [0.44573, 0.43259, 0.4301] # low noise
data[3][1] = [0.55761, 0.54497, 0.53731] # medium noise
data[3][2] = [0.6232, 0.59693, 0.591] # high noise

fig, axs = plt.subplots(2, 2, constrained_layout=True, sharey=True, sharex=True, figsize=(11,10))
for i in range(4):
    # axs[int(i%2)][int(i/2)].bar([0, 1, 2], data[i][0], fill=False)
    # axs[int(i%2)][int(i/2)].bar([4, 5, 6], data[i][1], fill=False, hatch='-')
    # axs[int(i%2)][int(i/2)].bar([8, 9, 10], data[i][2], fill=False, hatch='/')
    axs[int(i%2)][int(i/2)].plot([-1, 11], [no_noise_data[i]]*2, linestyle=':', color="Grey", label="no adding noise")
    axs[int(i%2)][int(i/2)].bar([0, 4, 8], [data[i][0][0], data[i][1][0], data[i][2][0]], fill=False, hatch='O.', label='with noise')
    axs[int(i%2)][int(i/2)].bar([1, 5, 9], [data[i][0][1], data[i][1][1], data[i][2][1]], fill=False, hatch='oo', label='1st denoise')
    axs[int(i%2)][int(i/2)].bar([2, 6, 10], [data[i][0][2], data[i][1][2], data[i][2][2]], fill=False, hatch='..', label='2nd denoise')
    axs[int(i%2)][int(i/2)].set_ylim(0.38, 0.72)
    axs[int(i%2)][int(i/2)].set_xlim(-1, 11)

for ax in axs.flat:
    ax.set_xticks([1, 5, 9], labels=['Low', 'Med', 'High'])

axs[1][0].set_xlabel('Noise level')
axs[1][1].set_xlabel('Noise level')
axs[0][0].set_ylabel('MAE (eV)')
axs[1][0].set_ylabel('MAE (eV)')


axs[0][0].text(0, 0.67, "(a) \'1by1\' on \nlinear noise")
axs[0][1].text(0, 0.67, "(b) \'1by1\' on \nsampled noise")
axs[1][0].text(0, 0.67, "(c) \'onion\' on \nlinear noise")
axs[1][1].text(0, 0.67, "(d) \'onion\' on \nsampled noise")
axs[0][0].legend(loc="upper right")
plt.savefig("Figure_3.png")
plt.show()